from typing import Dict, List, Optional

from instruction_re.core.datatypes import Instance, Preffix, Relation
from ..instances.EntityFormatter import Language
from ..format_utils import *
from ..instances.MainTaskFormatter import MainTaskFormatter


class QuestionFormatter:
    available_languages = [Language.EN.value]
    pattern = "What {1} {0}?"

    @classmethod
    def from_relations(
        cls, relations: Optional[List[Dict[str, List[str]]]], language: str = "en"
    ) -> Optional[str]:

        instrances = []
        for r in relations:
            instrances.append(cls.from_relation_triple(r, language))
        answer = ", ".join(instrances) + "."
        return answer

    @classmethod
    def from_relation_triple(
        cls, relation: Optional[Dict[str, List[str]]], language: str = "en"
    ) -> Optional[str]:
        if language not in cls.available_languages:
            raise ValueError(
                f"Expected language to be one of {cls.available_languages}"
            )

        if relation is None:
            return None

        return cls.pattern.format(relation["head"]["name"], relation["type"].lower())


class EntityPairExtractFormatter(MainTaskFormatter):
    """
    Task: Given sentence, subject and the relation type extract the object
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

    def format_data_sample(
        self,
        data_item,
        instruction: str,
    ) -> Instance:

        relations = data_item["relations"]

        instruction = add_instruction_prefix(instruction)
        conditions = add_conditions_prefix(self.format_type_answer(relations))
        sentence = add_sentence_prefix(data_item["sentence"])
        answer = super().format_answer(relations)

        return {
            "instruction": instruction,
            "context": sentence + " " + conditions + " " + Preffix.ANSWER.value,
            "labels": answer,
        }


def main():
    sample = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "relations": [
            {
                "head": {
                    "name": "John Wilkes Booth",
                    "type": "People",
                    "pos": [0, 17],
                },
                "type": "Kill",
                "tail": {
                    "name": "President Lincoln",
                    "type": "People",
                    "pos": [37, 54],
                },
            }
        ],
        "instruction": "Please extract all subjects according to input sentence, objects and relation types.",
    }

    formatter = EntityPairExtractFormatter()
    print(
        formatter.format_instance(
            context=sample["sentence"],
            relations=sample["relations"],
            instruction=sample["instruction"],
        )
    )


if __name__ == "__main__":
    main()
